
using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;

namespace Ack
{
	public class Huffman
	{
		private class Code
		{
			public byte _symbol;
			public BitArray _bits;
			public int _frequency;
			public byte _decoderCode;
			public byte _decoderCodeMask;

			public int BitCount
			{
				get
				{
					return _bits.Count;
				}
			}
			
			public override string ToString ()
			{
				string s = "";
				for(int bi = 0; bi < _bits.Count; ++bi) 
				{
					s += _bits.Get(bi) ? "1" : "0";
				}
				return s;
			}
		};
		
		private class HuffmanNode
		{
			public int _frequency;
			public bool _rightChild;
			public HuffmanNode _parent;
			public byte _symbol;
		
			public Code GetCode()
			{
				int codeLength = 0;
				HuffmanNode node = this;
				
				if(_frequency > 0)
				{
					while (node._parent != null) {
						codeLength++;
						node = node._parent;
					}
				}

				Code code = new Code();
				code._symbol = _symbol;
				code._frequency = _frequency;
				code._bits = new BitArray(codeLength);
				return code;
			}
		}

		private static void InsertSorted (LinkedList<HuffmanNode> nodes, HuffmanNode newNode)
		{
			if (nodes.Count == 0) {
				nodes.AddFirst (newNode);
			} else {
				LinkedListNode<HuffmanNode> largestSmaller = nodes.First;
				while (largestSmaller != null && largestSmaller.Value._frequency < newNode._frequency) {
					largestSmaller = largestSmaller.Next;
				}
				
				if (largestSmaller == null) {
					//Console.WriteLine("Adding f " + newNode._frequency + " last.");
					nodes.AddLast (newNode);
				} else {
					//Console.WriteLine("Adding f " + newNode._frequency + " before " + largestSmaller.Value._frequency);
					nodes.AddBefore (largestSmaller, newNode);
				}
			}
		}

		public static List<byte> Encode (List<byte> input, int symbolCount)
		{
			int[] symbolFreqs = new int[symbolCount];

			// Count symbol frequencies.
			for(int i = 0; i < input.Count; ++i)
			{
				symbolFreqs[input[i]]++;
			}
						
			// Make list of leaf nodes for the symbols and at the same time, make a reverse mapping from input value to each leaf.				
			HuffmanNode[] leaves = new HuffmanNode[symbolCount];
			for (int i = 0; i < symbolCount; ++i) 
			{
				leaves[i] = new HuffmanNode ();
				leaves[i]._symbol = (byte)i;
				leaves[i]._frequency = symbolFreqs[i];
			}
			
			// Insert leaves in sorted order in a list.
			LinkedList<HuffmanNode> nodes = new LinkedList<HuffmanNode> ();
			for (int i = 0; i < leaves.Length; ++i) {
				if(leaves[i]._frequency > 0)
					InsertSorted (nodes, leaves[i]);
			}
			
			// Build the tree by joining the two nodes with least frequency
			// as children of a new node with leftfreq+rightfreq.
			while (nodes.Count >= 2) {
				HuffmanNode lc = nodes.First.Value;
				nodes.RemoveFirst ();
				
				HuffmanNode rc = nodes.First.Value;
				nodes.RemoveFirst ();
				
				HuffmanNode newNode = new HuffmanNode ();
				newNode._frequency = lc._frequency + rc._frequency;
				lc._parent = newNode;
				lc._rightChild = false;
				rc._parent = newNode;
				rc._rightChild = true;
				
				InsertSorted (nodes, newNode);
			}
			
			// Generate code table
			Code[] codes = new Code[leaves.Length];
			for (int i = 0; i < leaves.Length; ++i) {
				codes[i] = leaves[i].GetCode();
			}
			
			MakeCodesCanonical (codes);
			
			/*for (int i = 0; i < codes.Length; ++i) {
				if(codes[i]._frequency>0)
					Console.WriteLine ("Freq for " + codes[i]._symbol + ": " + codes[i]._frequency + ", code: " + codes[i]);
			}*/
			
			// Count output bits (stupid, wouldn't be done if BitArray was growable)
			int outputBitCount = 0;
			for (int i = 0; i < input.Count; ++i) {
				
				int codeIndex = -1;
				for(int ci = 0; ci < codes.Length; ++ci) {
					if(codes[ci]._symbol == input[i])  {
						codeIndex = ci;
					}
				}
				
				outputBitCount += codes[codeIndex].BitCount;
			}
			
			List<byte> header = MakeHeader(codes);
			
			// Now translate all input symbols into codes
			int huffDataPos = 0;
			BitArray huffData = new BitArray(outputBitCount);
			for (int i = 0; i < input.Count; ++i) {
				
				int codeIndex = -1;
				for(int ci = 0; ci < codes.Length; ++ci) {
					if(codes[ci]._symbol == input[i])  {
						codeIndex = ci;
					}
				}
				
				Code code = codes[codeIndex];
				
				for (int j = 0; j < code._bits.Count; ++j)
					huffData[huffDataPos++] = code._bits.Get(j);
			}
			
			
			List<byte> outputBytes = new List<byte>();
			outputBytes.AddRange(BitConverter.GetBytes((ushort)(input.Count/2)));
			outputBytes.AddRange(header);
					
			for(int i = 0; i < (huffData.Count/8+1); ++i)
			{
				outputBytes.Add( (byte) (
					((i*8+0) < huffData.Count ? (huffData[i*8+0] ? 0x01 : 0) : 0) |
					((i*8+1) < huffData.Count ? (huffData[i*8+1] ? 0x02 : 0) : 0) |
					((i*8+2) < huffData.Count ? (huffData[i*8+2] ? 0x04 : 0) : 0) |
					((i*8+3) < huffData.Count ? (huffData[i*8+3] ? 0x08 : 0) : 0) |
					((i*8+4) < huffData.Count ? (huffData[i*8+4] ? 0x10 : 0) : 0) |
					((i*8+5) < huffData.Count ? (huffData[i*8+5] ? 0x20 : 0) : 0) |
					((i*8+6) < huffData.Count ? (huffData[i*8+6] ? 0x40 : 0) : 0) |
					((i*8+7) < huffData.Count ? (huffData[i*8+7] ? 0x80 : 0) : 0) ) );
			}
			
			return outputBytes;
		}

		private static List<byte> MakeHeader(Code[] codes)
		{
			// Codes must be canonical and symbol value ordered.
			
			List<byte> codeLengths = new List<byte>();
			
			int maxCodeLength = 0;
			for(int ci = 0; ci < codes.Length; ++ci)
			{
				int len = codes[ci].BitCount;
				Debug.Assert(len < 16);
				codeLengths.Add((byte)len);
				if(len > maxCodeLength)
					maxCodeLength = len;
			}

			List<byte> cmp = new List<byte>(3*16);
			for(int i = 0; i < 16; ++i)
			{
				cmp.Add(codes[i]._decoderCode);
			}
			for(int i = 0; i < 16; ++i)
			{
				cmp.Add(codes[i]._decoderCodeMask);
			}
			for(int i = 0; i < 16; ++i)
			{
				cmp.Add((byte)codes[i].BitCount);
			}
						return cmp;
			/*			
			List<byte> cmp = new List<byte>();
			for(int i = 0; i < codes.Length/2; ++i)
			{
				cmp.Add((byte)((codeLengths[i*2+1]<<4)+codeLengths[i*2]));
			}
			
		//	Console.WriteLine("Huffman: Compressed header bytes: " + cmp.Count);
			return cmp;
*/		}
		
		// Make the codes "canonical", which allows less header data to be required
		// because the decoder can reconstruct the codes knowing just the lengths.
		// See: http://en.wikipedia.org/wiki/Canonical_Huffman_code
		private static void MakeCodesCanonical(Code[] codes)
		{
			// Sort the codes by ascending length
			Array.Sort(codes, delegate(Code a, Code b) 
			                      {
				int c = a.BitCount.CompareTo (b.BitCount);

				if(c==0)
				{
					c = a._symbol.CompareTo(b._symbol);
				}
				
				return c;
			});

			// Generate the codes.
			int currentCodeLength = codes[0].BitCount;
			uint currentCode = 0;
			for(int ci = 0; ci < codes.Length; ++ci)
			{
				BitArray bits = codes[ci]._bits;
				bits.SetAll(false);

				for(int bi = 0; bi < currentCodeLength; ++bi)
				{
					bits.Set(currentCodeLength-bi-1, (currentCode & (1<<bi)) != 0);
				}
				
				codes[ci]._decoderCode = HuffmanDecoder.GenerateCodeEntry((int)currentCode,currentCodeLength);
				codes[ci]._decoderCodeMask = HuffmanDecoder.GenerateMaskEntry(currentCodeLength);
				
				currentCode++;
				
				if(ci < codes.Length-1)
				{
					int longer = codes[ci+1].BitCount - bits.Count;
					currentCode <<= longer;
					currentCodeLength += longer;
				}
			}
			
			// Sort the codes back into symbol-value order so we can look them up rapidly.
			Array.Sort(codes, delegate(Code a, Code b) {
				return a._symbol.CompareTo(b._symbol);
			});
		}
	}

}
